#include "framebuffer.h"
#include <cassert>
#include <GL/glew.h>
#include "stb_image.h"
#include "stb_image_write.h"

namespace kanon
{
    FrameBuffer::FrameBuffer()
    {
        fbo = 0;
        rbo = 0;
        color_buffers.clear();
        width = 0;
        height = 0;
    }

    bool FrameBuffer::Init(uint32_t width, uint32_t height, uint32_t num_color_buffers, bool depthstencil)
    {
        this->width = width;
        this->height = height;

        glGenFramebuffers(1, &fbo);
        glBindFramebuffer(GL_FRAMEBUFFER, fbo);

        if (num_color_buffers <= 0)
            return false;

        for (uint32_t i = 0; i < num_color_buffers; ++i)
        {
            uint32_t texture;
            glGenTextures(1, &texture);
            glBindTexture(GL_TEXTURE_2D, texture);
            glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height,
                0, GL_RGB, GL_UNSIGNED_BYTE, NULL);
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
            glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, texture, 0);
            color_buffers.push_back(texture);
        }

        // create depth stencil
        glGenRenderbuffers(1, &rbo);
        glBindRenderbuffer(GL_RENDERBUFFER, rbo);
        glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, width, height);
        glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT,
            GL_RENDERBUFFER, rbo);

        // check status
        if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
            return false;

        glBindFramebuffer(GL_FRAMEBUFFER, 0);

        return true;
    }

    bool FrameBuffer::Init(uint32_t width, uint32_t height, const std::shared_ptr<Texture2D>& texture)
    {
        this->width = width;
        this->height = height;

        glGenFramebuffers(1, &fbo);
        glBindFramebuffer(GL_FRAMEBUFFER, fbo);

        uint32_t texture_id = texture->getId();
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, texture_id, 0);
        color_buffers.push_back(texture_id);

        return true;
    }

    void FrameBuffer::Bind() const
    {
        glBindFramebuffer(GL_FRAMEBUFFER, fbo);
    }

    void FrameBuffer::Bind(const std::vector<uint32_t>& buffer_indices) const
    {
        glBindFramebuffer(GL_FRAMEBUFFER, fbo);
        std::vector<GLenum> attachments;
        for (auto& idx : buffer_indices)
            attachments.push_back(GL_COLOR_ATTACHMENT0 + idx);
        glDrawBuffers(buffer_indices.size(), attachments.data());
    }

    void FrameBuffer::Unbind() const
    {
        glBindFramebuffer(GL_FRAMEBUFFER, 0);
    }

    void FrameBuffer::Capture(uint32_t index, const char* filepath)
    {
        // 
        std::vector<unsigned char> pixels(this->width * this->height * 4);
        glBindTexture(GL_TEXTURE_2D,this->color_buffers[index]);
        glGetTexImage(GL_TEXTURE_2D, 0, GL_RGBA, GL_UNSIGNED_BYTE, pixels.data());

        stbi_write_png(filepath, this->width, this->height, 4, pixels.data(), this->width * 4);
    }

    void FrameBuffer::CaptureTexture(uint32_t texture_id, uint32_t width, uint32_t height, const char* filepath)
    {
        std::vector<unsigned char> pixels(width * height * 4);
        glBindTexture(GL_TEXTURE_2D, texture_id);
        glGetTexImage(GL_TEXTURE_2D, 0, GL_RGBA, GL_UNSIGNED_BYTE, pixels.data());

        stbi_write_png(filepath, width, height, 4, pixels.data(), width * 4);
    }
}
